import torch
from torch.nn import functional as f
from d2l import torch as d2l
import text_preprocess as tp

batch_size = 32
num_steps = 35
train_iter,vocab = tp.load_data_time_machine(batch_size,num_steps)

def get_params(vocab_size,num_hiddens,device):
    num_inputs = num_outputs = vocab_size
    def normal(shape):
        return torch.randn(size=shape,device=device) * 0.01
    W_x_hidden = normal((num_inputs,num_hiddens))
    W_hidden_hidden = normal((num_hiddens,num_hiddens))
    b_hidden = torch.zeros(num_hiddens,device=device)
    W_hidden_o = normal((num_hiddens,num_outputs))
    b_o = torch.zeros(num_outputs,device=device)
    params = [W_x_hidden,W_hidden_hidden,b_hidden,W_hidden_o,b_o]
    for param in params:
        param.requires_grad_(True)
    return params

def init_rnn_state(batch_size,num_hiddens,device):
    return (torch.zeros((batch_size,num_hiddens),device=device),)

def rnn(inputs,state,params):
    W_xh,W_hh,b_h,W_ho,b_o = params
    H, = state
    outputs = []
    for X in inputs:
        H = torch.tanh(torch.mm(X,W_xh) + torch.mm(H,W_hh) + b_h)
        Y = torch.mm(H, W_ho) + b_o
        outputs.append(Y)
    return torch.cat(outputs,dim=0), (H,)

class RNNModelScratch:
    def __init__(self,vocab_size,num_hiddens,device,get_params,init_state,forward_fn):
        self.vocab_size = vocab_size
        self.num_hiddens = num_hiddens
        self.params = get_params(vocab_size,num_hiddens,device)
        self.init_state = init_state
        self.forward_fn = forward_fn
    
    def __call__(self,X,state):
        X = f.one_hot(X.T,self.vocab_size).type(torch.float32)
        return self.forward_fn(X,state,self.params)
    
    def begin_state(self,batch_size,device):
        return self.init_state(batch_size,self.num_hiddens,device)
    
num_epochs = 500
lr = 1
num_hiddens = 512
net = RNNModelScratch(len(vocab),num_hiddens,d2l.try_gpu(),get_params,init_rnn_state,rnn)
import RNN_train as rt
rt.train(net,train_iter,vocab,lr,num_epochs,d2l.try_gpu(),use_random_iter=False)


    
